from fastapi import APIRouter, HTTPException, UploadFile, File, Form
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
from rag import RAGSystem
from monitoring.decorators import monitor_api_endpoint, monitor_rag_query, monitor_document_upload
import os
import tempfile
import logging
from config.thresholds import RAGThresholds, ValidationThresholds

logger = logging.getLogger(__name__)

# 创建路由器
router = APIRouter()

# 初始化RAG系统
rag_system = RAGSystem()

# 请求和响应模型
class QuestionRequest(BaseModel):
    question: str
    collection_name: Optional[str] = "default"
    top_k: Optional[int] = RAGThresholds.DEFAULT_TOP_K
    rerank_threshold: Optional[float] = RAGThresholds.DEFAULT_RERANK_THRESHOLD

class QuestionResponse(BaseModel):
    status: str
    question: str
    answer: str
    sources: List[Dict[str, Any]]
    context_length: int
    documents_retrieved: int

class UploadResponse(BaseModel):
    status: str
    message: str
    collection_info: Dict[str, Any]
    file_path: str

class CollectionInfoResponse(BaseModel):
    collection_name: str
    document_count: int
    embedding_model: str

# 文档上传
@router.post("/upload", response_model=UploadResponse)
@monitor_api_endpoint("rag_upload_document")
@monitor_document_upload()
async def upload_document(
    file: UploadFile = File(...),
    collection_name: str = Form("default")
):
    """上传文档到RAG系统"""
    try:
        logger.info(f"收到文档上传请求: {file.filename}")
        
        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as temp_file:
            # 写入文件内容
            content = await file.read()
            temp_file.write(content)
            temp_file_path = temp_file.name
        
        # 处理文档
        result = rag_system.upload_document(temp_file_path, collection_name)
        
        # 删除临时文件
        os.unlink(temp_file_path)
        
        if result["status"] == "success":
            return UploadResponse(**result)
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"文档上传失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"文档上传失败: {str(e)}")

# 问答
@router.post("/ask", response_model=QuestionResponse)
@monitor_api_endpoint("rag_ask_question")
@monitor_rag_query()
async def ask_question(request: QuestionRequest):
    """基于RAG系统回答问题"""
    try:
        logger.info(f"收到问题: {request.question}")
        
        result = rag_system.ask_question(
            question=request.question,
            collection_name=request.collection_name,
            top_k=request.top_k,
            rerank_threshold=request.rerank_threshold
        )
        
        if result["status"] == "success":
            return QuestionResponse(**result)
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"回答问题失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"回答问题失败: {str(e)}")

# 获取集合信息
@router.get("/collection/{collection_name}", response_model=CollectionInfoResponse)
@monitor_api_endpoint("rag_get_collection_info")
async def get_collection_info(collection_name: str = "default"):
    """获取集合信息"""
    try:
        result = rag_system.get_collection_info(collection_name)
        
        if "status" in result and result["status"] == "error":
            raise HTTPException(status_code=400, detail=result["message"])
        
        return CollectionInfoResponse(**result)
        
    except Exception as e:
        logger.error(f"获取集合信息失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取集合信息失败: {str(e)}")

# 列出所有集合
@router.get("/collections")
@monitor_api_endpoint("rag_list_collections")
async def list_collections():
    """列出所有集合"""
    try:
        collections = rag_system.list_collections()
        return {
            "collections": collections,
            "count": len(collections)
        }
        
    except Exception as e:
        logger.error(f"列出集合失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"列出集合失败: {str(e)}")

# 删除集合
@router.delete("/collection/{collection_name}")
@monitor_api_endpoint("rag_delete_collection")
async def delete_collection(collection_name: str):
    """删除集合"""
    try:
        result = rag_system.delete_collection(collection_name)
        
        if result["status"] == "success":
            return result
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"删除集合失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"删除集合失败: {str(e)}")

# 健康检查
@router.get("/health")
@monitor_api_endpoint("rag_health_check")
async def health_check():
    """RAG系统健康检查"""
    return {
        "status": "healthy",
        "service": "RAG System",
        "features": ["document_upload", "question_answering", "vector_search"]
    } 